from openai import OpenAI
import os
import base64
from .prompt import image_system_prompt, user_prompt
from src.utils import get_logger
import json

logger = get_logger(__name__)

#  base 64 编码格式
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def get_image_response(image_path, history_messages=None):
    """
    获取模型响应，并使用滑动窗口管理历史记录。

    Args:
        image_path: 图片路径。
        history_messages: 历史消息列表，用于维护上下文。默认为 None。

    Returns:
        tuple: 包含模型响应内容和更新后的历史消息列表。
    """
    base64_image = encode_image(image_path)
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        base_url=os.getenv("DASHSCOPE_BASE_URL"),
    )

    # 初始消息列表，包含 system_prompt
    if history_messages is None:
        messages = [
            {
                "role": "system",
                "content": image_system_prompt
            }
        ]
    else:
        messages = history_messages

    # 添加用户消息
    user_message = {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": "分析画面"
            },
            {
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{base64_image}"
                }
            }
        ]
    }
    messages.append(user_message)

    completion = client.chat.completions.create(
        model="qwen-vl-plus",
        messages=messages,
    )
    response_content = completion.choices[0].message.content
    logger.info(f"模型回复: \n{response_content}")
    
    # 将响应内容转换为字典格式
    try:
        # 尝试将响应解析为JSON
        response_dict = json.loads(response_content)
    except json.JSONDecodeError:
        # 如果无法解析为JSON，则创建一个包含原始文本的字典
        response_dict = {
            "description": response_content,
            "timestamp": os.path.basename(image_path).split("_")[1].split(".")[0]
        }
    
    # 添加模型回复到历史消息
    assistant_message = {
        "role": "assistant",
        "content": response_content
    }
    messages.append(assistant_message)

    # 滑动窗口逻辑：保留 system_prompt 和最近 5 轮对话 (user + assistant)
    if len(messages) > 11: # system prompt + 5 user + 5 assistant = 11
        messages = [messages[0]] + messages[-10:] # 保留 system_prompt 和最近 10 条消息 (5轮对话)

    return response_dict, messages # 返回字典格式的响应和更新后的消息列表
                    

if __name__=='__main__':
    history = None # 初始历史消息为空
    # image_paths = ["/root/autodl-tmp/code/agent/src/file/extract_picture/456998289-1-192/keyframe_00:01.200.jpg",
    #                "/root/autodl-tmp/code/agent/src/file/extract_picture/456998289-1-192/keyframe_00:05.440.jpg",
    #                "/root/autodl-tmp/code/agent/src/file/extract_picture/456998289-1-192/keyframe_00:09.280.jpg",
    #                "/root/autodl-tmp/code/agent/src/file/extract_picture/456998289-1-192/keyframe_00:13.040.jpg"] # 假设有多个图片路径

    image_paths = ["/root/autodl-tmp/code/agent/src/file/extract_picture/456998289-1-192/keyframe_00:34.160.jpg",] 
    for image_path in image_paths:
        res, history = get_response(image_path, history) # 传递历史消息
        print(f"图片: {image_path}")
        print(f"模型回复: {res}")
        print(f"当前历史消息长度: {len(history) if history else 0}")
        print("-" * 20)